# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

from absl import app, flags

import collections
import functools
import logging
import math

import pandas
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from coin.base.datetime_util import to_timestamp_int
from coin.base.book.types import TradeSide

import coin.strategy.mm.tool.archive_base as abase


class DetectSweep(object):
  def __init__(self, product, start_ts=None, end_ts=None):
    self._product = product
    self._start_ts = to_timestamp_int(start_ts)
    self._end_ts = to_timestamp_int(end_ts or 0)

    self._bidarr = None
    self._askarr = None

    self._bqty_map = collections.defaultdict(float)
    self._sqty_map = collections.defaultdict(float)

    self._prev_trade = None

    self._last_sweep = {}
    self._stats = []
    self._sweepid = 1

  def on_book_reset(self, book_builder_name, book_builder):
    for product in [self._product]:
      book_builder.subscribe(product, functools.partial(self.on_book, product))
      if hasattr(book_builder, "subscribe_trade"):
        book_builder.subscribe_trade(product, functools.partial(self.on_trade, product))

  def pprint(self, *args):
    pass
    # print(*args)

  def on_book(self, product, book):
    if book.timestamp < self._start_ts:
      return

    if self._end_ts > 0 and book.timestamp > self._end_ts:
      return

    try:
      bid0_price = book.bid0().price
      ask0_price = book.ask0().price
    except (AttributeError, IndexError):
      return

    bidarr = book.get_bid_array(1)
    askarr = book.get_ask_array(1)

    if (self._bidarr is not None and self._askarr is not None):
      if (self._bidarr[0, 0] > bidarr[0, 0] or self._askarr[0, 0] < askarr[0, 0]):
        if self._bidarr[0, 0] > bidarr[0, 0]:
          self.pprint("DOWN TICK")
          side = "SELL"
        else:
          self.pprint("UP TICK")
          side = "BUY"
        all_none = True
        for sweeptype in ["price", "qty"]:
          strade = self._last_sweep.get((side, sweeptype), None)
          if strade is not None:
            tdms = ((book.timestamp - strade[2]) * 1e-6)
            self._stats.append((self._sweepid, sweeptype, tdms))
            all_none = False
          else:
            tdms = None
          self.pprint(tdms, side, sweeptype, strade)
        self._stats.append((self._sweepid, "book", 0))

        self.pprint(self._bidarr, self._askarr)
        self._last_sweep.clear()
        self._sweepid += 1
        self.pprint(bidarr, askarr)

    self._bidarr = book.get_bid_array(1)
    self._askarr = book.get_ask_array(1)
    self._bqty_map.clear()
    self._sqty_map.clear()

  def add_sweep(self, side, sweeptype, trade):
    if (side, sweeptype) not in self._last_sweep:
      self._last_sweep[(side,
                        sweeptype)] = (side, sweeptype, trade.timestamp, trade.price, trade.qty)

  def on_trade(self, product, trade):
    exchange = product.exchange
    if trade.timestamp < self._start_ts:
      return

    if self._end_ts > 0 and trade.timestamp > self._end_ts:
      return

    self._prev_trade = trade

    if trade.side == TradeSide.TRADE_BUY_SIDE:
      self._bqty_map[trade.price] += trade.qty
      if self._askarr is not None:
        if (trade.price > self._askarr[0, 0]
            or (math.isclose(trade.price, self._askarr[0, 0]) and trade.qty >= self._askarr[0, 1])):
          sweeptype = ("price" if trade.price > self._askarr[0, 0] else "qty")
          self.add_sweep("BUY", sweeptype, trade)
    elif trade.side == TradeSide.TRADE_SELL_SIDE:
      self._sqty_map[trade.price] += trade.qty
      if self._bidarr is not None:
        if (trade.price < self._bidarr[0, 0]
            or (math.isclose(trade.price, self._bidarr[0, 0]) and trade.qty >= self._bidarr[0, 1])):
          sweeptype = ("price" if trade.price < self._bidarr[0, 0] else "qty")
          self.add_sweep("SELL", sweeptype, trade)

  def summary(self):
    df = pandas.DataFrame(self._stats, columns=['id', 'method', 'tdms'])

    n = len(df.groupby('method'))

    coverage = (df.groupby('method').count() / df.groupby('method').count().max())['id']

    plt.suptitle(flags.FLAGS.trading_date)
    for i, (method, grpdf) in enumerate(df.groupby('method')):
      tdmssrt = sorted(grpdf['tdms'])
      plt.subplot(n, 4, (i * 4) + 1)
      plt.plot(tdmssrt, linewidth=0.5)
      plt.title("%s info : covers %.0f%%" % (method, coverage[method] * 100))

      statser = pandas.Series(tdmssrt).describe().astype(int)
      plt.subplot(n, 4, (i * 4) + 2)
      plt.table(cellText=statser.to_frame().reset_index().values, loc="center")
      plt.title("time(ms) advantage to book feed")
      plt.axis("off")

      tdmscut = tdmssrt[:int(len(tdmssrt) * 0.8)]
      plt.subplot(n, 4, (i * 4) + 3)
      plt.plot(tdmscut, linewidth=0.5)
      plt.title("cutting [0,80%] outlier")
      plt.subplot(n, 4, (i * 4) + 4)
      plt.hist(tdmscut, bins=100)
      plt.title("cutting [0,80%] outlier hist")
    plt.tight_layout()
    plt.savefig("pic/detectsweep_%s.png" % flags.FLAGS.trading_date)
    plt.close()


def launch(baseflags, fsflags, symbols, out_dir):
  plt.rcParams['figure.figsize'] = 12, 8
  plt.rcParams['font.size'] = 10
  plt.rcParams['legend.fontsize'] = 10
  plt.rcParams['xtick.labelsize'] = 10
  plt.rcParams['ytick.labelsize'] = 10

  products = abase.symbols_to_products(symbols, baseflags)
  start_time, end_time = abase.get_time_range(baseflags)

  assert len(products) == 1

  strategy = DetectSweep(product=products[0], start_ts=start_time, end_ts=end_time)

  abase.run_from_archive_base(baseflags=baseflags,
                              feed_sub_flags_proto=fsflags,
                              products=products,
                              on_book_reset=strategy.on_book_reset)

  strategy.summary()


def main(argv):
  baseflags = abase.get_base_flags()
  fsflags = abase.get_feed_subscription_flags()

  symbols = flags.FLAGS.symbol.split(",")
  out_dir = flags.FLAGS.out_dir
  launch(baseflags, fsflags, symbols, out_dir)


'''
./pyrunner coin/strategy/mm/tool/detect_sweep.py \
--trading_date=20190620 \
--exchange=Bitmex --market_type=Futures \
--symbol=BTC-USD.PERPETUAL \
--feed_machine=feed-01.ap-northeast-1.aws \
--use_feed_cache --time_range=0-2
'''

if __name__ == '__main__':
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  # fill up default settings
  flags.FLAGS.feed_machine = 'feed-01.ap-northeast-1.aws'
  flags.DEFINE_string('symbol', None, 'symbol name')
  flags.DEFINE_string('out_dir', 'pic', 'Output directory.')
  flags.FLAGS.use_feed_cache = True

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')

  app.run(main)
